package com.vclee.fast.integration.imserver;

import com.vclee.fast.utils.model.Bson;
import lombok.SneakyThrows;
import okhttp3.*;
import okhttp3.sse.EventSource;
import okhttp3.sse.EventSourceListener;
import okhttp3.sse.EventSources;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Component;

import javax.net.ssl.SSLContext;
import javax.net.ssl.TrustManager;
import javax.net.ssl.TrustManagerFactory;
import javax.net.ssl.X509TrustManager;
import java.security.KeyStore;
import java.security.SecureRandom;
import java.security.cert.X509Certificate;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;

@Component
public class XunFeiService {

    private static final Logger logger = LoggerFactory.getLogger(XunFeiService.class);

    private final static String URL = "https://spark-api-open.xf-yun.com/v1/chat/completions";
    private final static String KEY = "YvJwYfHEMJcJCjhGQEhT:MeoJAgLcIJKljzdwdPFT";
    private final static String MODEL = "generalv3.5"; // https://www.xfyun.cn/doc/spark/Web.html#_1-%E6%8E%A5%E5%8F%A3%E8%AF%B4%E6%98%8E
    private final static Map<Long, List<Bson>> messages = new ConcurrentHashMap<>();


    @SneakyThrows
    public static void conversation(Long userId, Long cvsId, String text, Callback<String> callback) {
        Bson message = new Bson();
        message.put("role", "user");
        message.put("content", text);

        StringBuilder words = new StringBuilder();
        List<Bson> list = messages.get(userId);
        if (list == null) list = new ArrayList<>();
        list.add(message);
        if (list.size() > 10) {
            list.remove(0);
        }
        messages.put(userId, list);

        Bson params = new Bson();
        params.put("model", MODEL);
        params.put("messages", list);
//        params.put("temperature", 0);
//        params.put("max_tokens", 2048);
        params.put("stream", true);

        logger.info("数组请求参数:{}", params);

        Request request = new Request.Builder()
                .url(URL)
                .addHeader("Content-Type", "application/json")
                .addHeader("Access-Control-Allow-Origin", "*")
                .addHeader("Access-Control-Allow-Methods", "*")
                .addHeader("Access-Control-Allow-Headers", "Content-Type,XFILENAME,XFILECATEGORY,XFILESIZE")
                .addHeader("Authorization", "Bearer " + KEY)
                .addHeader("Accept", "text/event-stream")
                .post(RequestBody.create(MediaType.parse("application/json"), params.toString()))
                .build();

        OkHttpClient.Builder builder = new OkHttpClient.Builder();
        builder.connectTimeout(600, TimeUnit.SECONDS);
        builder.writeTimeout(600, TimeUnit.SECONDS);
        builder.readTimeout(600, TimeUnit.SECONDS);
        SSLContext ctx = SSLContext.getInstance("SSL");
        ctx.init(null, new TrustManager[]{new X509TrustManager() {
            @Override
            public void checkClientTrusted(X509Certificate[] chain, String authType) {
            }

            @Override
            public void checkServerTrusted(X509Certificate[] chain, String authType) {
            }

            @Override
            public X509Certificate[] getAcceptedIssuers() {
                return new X509Certificate[]{};
            }
        }}, new SecureRandom());

        TrustManagerFactory trustManagerFactory = TrustManagerFactory.getInstance(TrustManagerFactory.getDefaultAlgorithm());
        trustManagerFactory.init((KeyStore) null);
        X509TrustManager trustManager = (X509TrustManager) trustManagerFactory.getTrustManagers()[0];

        builder.sslSocketFactory(ctx.getSocketFactory(), trustManager)
                .hostnameVerifier((s, sslSession) -> true);
        OkHttpClient client = builder.build();

        EventSource.Factory factory = EventSources.createFactory(client);
        factory.newEventSource(request, new EventSourceListener() {
            @Override
            @SuppressWarnings("all")
            public void onEvent(EventSource eventSource, String id, String type, String data) {
                // logger.info("AsyncClient.onEvent, id={}, type={}, data= {}", id, type, data);
                Bson res = new Bson();
                res.put("cvsId", cvsId);
                if (data.equalsIgnoreCase("[DONE]")) {
                    res.put("finish", true);
                    Bson item = new Bson();
                    item.put("role", "assistant");
                    item.put("content", words);
                    messages.get(userId).add(item);
                } else {
                    res.put("finish", false);
                    Bson json = Bson.create(data).getBsonArray("choices").get(0).getBson("delta");
                    words.append(json.getString("content"));
                    res.put("content", json.getString("content"));
                }
                callback.call(res.toString());
            }

            @Override
            @SuppressWarnings("all")
            public void onFailure(EventSource eventSource, Throwable t, Response response) {
                Bson json = new Bson();
                json.put("finish", true);
                json.put("content", "请重新表述一下你的问题");
                callback.call(json.toString());
                logger.error("asyncClient.execute failed:{}, e: ", t.getMessage(), t);
                eventSource.cancel();
            }
        });
    }
}
